package cn.coder.toolkit;

import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.PSSParameterSpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA 加密工具
 * <p>
 * 公钥加密/私钥解密
 * 私钥签名/公钥验签
 * 密钥生成/密钥转换
 * <p>
 * 注意: JAVA 生成的私钥是 PKCS8 格式, 公钥是 X.509 格式
 * <p>
 * 业内规范: 字符串格式的 密钥/密文/签名 等默认都使用 BASE64 编码
 */
@Slf4j
public final class RSAKit {

	private RSAKit() {}

	/**
	 * 加密算法
	 */
	private static final String RSA_ALGORITHM = "RSA";

	@Getter
	@AllArgsConstructor
	public enum KeySizeEnum {
		SIZE_2048(2048),
		SIZE_4096(4096);
		private final int length;
	}

	@Getter
	@AllArgsConstructor
	public enum CipherAlgorithmEnum {
		RSA_OAEP_SHA256("RSA/ECB/OAEPWithSHA-256AndMGF1Padding", new OAEPParameterSpec("SHA-256", "MGF1", MGF1ParameterSpec.SHA256, PSource.PSpecified.DEFAULT)),
		RSA_OAEP_SHA512("RSA/ECB/OAEPWithSHA-512AndMGF1Padding", new OAEPParameterSpec("SHA-512", "MGF1", MGF1ParameterSpec.SHA512, PSource.PSpecified.DEFAULT));
		private final String algorithm;
		private final OAEPParameterSpec oaepParameterSpec;
	}

	// 以下是密钥相关方法 ---------- ---------- ---------- ---------- ----------

	/**
	 * 生成公私钥对
	 *
	 * @param keySize 密钥长度, 2048/4096
	 */
	public static KeyPair generateKeyPair(KeySizeEnum keySize) {
		try {
			KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(RSA_ALGORITHM);
			keyPairGenerator.initialize(keySize.getLength());
			return keyPairGenerator.generateKeyPair();
		} catch (Throwable cause) {
			throw new RuntimeException(cause);
		}
	}

	public static KeyPair generateKeyPair() {
		return generateKeyPair(KeySizeEnum.SIZE_2048);
	}

	/**
	 * 字符串公钥转标准公钥
	 */
	public static PublicKey toPublicKey(String publicKeyStr) {
		try {
			return KeyFactory.getInstance(RSA_ALGORITHM).generatePublic(new X509EncodedKeySpec(Base64.getDecoder().decode(publicKeyStr)));
		} catch (Throwable cause) {
			throw new RuntimeException(cause);
		}
	}

	/**
	 * 字符串私钥转标准私钥
	 */
	public static PrivateKey toPrivateKey(String privateKeyStr) {
		try {
			return KeyFactory.getInstance(RSA_ALGORITHM).generatePrivate(new PKCS8EncodedKeySpec(Base64.getDecoder().decode(privateKeyStr)));
		} catch (Throwable cause) {
			throw new RuntimeException(cause);
		}
	}

	/**
	 * 标准密钥转字符串密钥
	 */
	public static String toKeyStr(Key key) {
		return Base64.getEncoder().encodeToString(key.getEncoded());
	}

	// 以下是密钥加密解密通用依赖方法 ---------- ---------- ---------- ---------- ----------

	/**
	 * 计算最大分段长度，keySize单位bit
	 */
	private static int getMaxEncryptBlockSize(CipherAlgorithmEnum cipherAlgorithmEnum, KeySizeEnum keySize) {
		// RSA 最大明文长度 = keySize_in_bytes - 2 * hash_len - 2 （OAEP填充公式）
		int hashLen;
		switch (cipherAlgorithmEnum) {
			case RSA_OAEP_SHA256:
				hashLen = 32; // SHA-256 输出长度32字节
				break;
			case RSA_OAEP_SHA512:
				hashLen = 64; // SHA-512 输出长度64字节
				break;
			default:
				throw new RuntimeException("不支持的 Cipher Algorithm");
		}
		return keySize.getLength() / 8 - 2 * hashLen - 2;
	}

	/**
	 * 分段加密
	 */
	public static byte[] encryptByPublicKey(PublicKey publicKey, CipherAlgorithmEnum cipherAlgorithm, KeySizeEnum keySize, byte[] data) {
		try {
			Cipher cipher = Cipher.getInstance(cipherAlgorithm.getAlgorithm());
			cipher.init(Cipher.ENCRYPT_MODE, publicKey, cipherAlgorithm.getOaepParameterSpec());
			int maxBlock = getMaxEncryptBlockSize(cipherAlgorithm, keySize);
			try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
				for (int offset = 0; offset < data.length; offset += maxBlock) {
					int length = Math.min(maxBlock, data.length - offset);
					byte[] encryptedBlock = cipher.doFinal(data, offset, length);
					out.write(encryptedBlock);
				}
				return out.toByteArray();
			}
		} catch (Exception e) {
			throw new RuntimeException("分段加密失败", e);
		}
	}

	/**
	 * 分段解密，密钥长度/8 是每个密文块大小
	 */
	public static byte[] decryptByPrivateKey(PrivateKey privateKey, CipherAlgorithmEnum cipherAlgorithm, KeySizeEnum keySize, byte[] encryptedData) {
		try {
			Cipher cipher = Cipher.getInstance(cipherAlgorithm.getAlgorithm());
			cipher.init(Cipher.DECRYPT_MODE, privateKey, cipherAlgorithm.getOaepParameterSpec());
			int cipherBlockSize = keySize.getLength() / 8;
			try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
				for (int offset = 0; offset < encryptedData.length; offset += cipherBlockSize) {
					int length = Math.min(cipherBlockSize, encryptedData.length - offset);
					byte[] decryptedBlock = cipher.doFinal(encryptedData, offset, length);
					out.write(decryptedBlock);
				}
				return out.toByteArray();
			}
		} catch (Exception e) {
			throw new RuntimeException("分段解密失败", e);
		}
	}

	// 以下是公钥加密相关方法 ---------- ---------- ---------- ---------- ----------

	/**
	 * 公钥加密
	 */
	public static byte[] encryptByPublicKey(PublicKey publicKey, byte[] data) {
		return encryptByPublicKey(publicKey, CipherAlgorithmEnum.RSA_OAEP_SHA256, KeySizeEnum.SIZE_2048, data);
	}
	public static byte[] encryptByPublicKey(PublicKey publicKey, String data) {
		return encryptByPublicKey(publicKey, data.getBytes(StandardCharsets.UTF_8));
	}
	public static byte[] encryptByPublicKeyStr(String publicKeyStr, byte[] data) {
		return encryptByPublicKey(toPublicKey(publicKeyStr), data);
	}
	public static byte[] encryptByPublicKeyStr(String publicKeyStr, String data) {
		return encryptByPublicKey(toPublicKey(publicKeyStr), data.getBytes(StandardCharsets.UTF_8));
	}
	public static String encryptToStrByPublicKey(PublicKey publicKey, byte[] data) {
		return Base64.getEncoder().encodeToString(encryptByPublicKey(publicKey, data));
	}
	public static String encryptToStrByPublicKey(PublicKey publicKey, String data) {
		return Base64.getEncoder().encodeToString(encryptByPublicKey(publicKey, data.getBytes(StandardCharsets.UTF_8)));
	}
	public static String encryptToStrByPublicKeyStr(String publicKeyStr, byte[] data) {
		return Base64.getEncoder().encodeToString(encryptByPublicKey(toPublicKey(publicKeyStr), data));
	}
	public static String encryptToStrByPublicKeyStr(String publicKeyStr, String data) {
		return Base64.getEncoder().encodeToString(encryptByPublicKey(toPublicKey(publicKeyStr), data.getBytes(StandardCharsets.UTF_8)));
	}

	// 以下是私钥解密相关方法 ---------- ---------- ---------- ---------- ----------

	/**
	 * 私钥解密
	 */
	public static byte[] decryptByPrivateKey(PrivateKey privateKey, byte[] data) {
		return decryptByPrivateKey(privateKey, CipherAlgorithmEnum.RSA_OAEP_SHA256, KeySizeEnum.SIZE_2048, data);
	}
	public static byte[] decryptByPrivateKey(PrivateKey privateKey, String data) {
		return decryptByPrivateKey(privateKey, Base64.getDecoder().decode(data));
	}
	public static byte[] decryptByPrivateKeyStr(String privateKeyStr, byte[] data) {
		return decryptByPrivateKey(toPrivateKey(privateKeyStr), data);
	}
	public static byte[] decryptByPrivateKeyStr(String privateKeyStr, String data) {
		return decryptByPrivateKey(toPrivateKey(privateKeyStr), Base64.getDecoder().decode(data));
	}
	public static String decryptToStrByPrivateKey(PrivateKey privateKey, byte[] data) {
		return new String(decryptByPrivateKey(privateKey, data), StandardCharsets.UTF_8);
	}
	public static String decryptToStrByPrivateKey(PrivateKey privateKey, String data) {
		return new String(decryptByPrivateKey(privateKey, Base64.getDecoder().decode(data)), StandardCharsets.UTF_8);
	}
	public static String decryptToStrByPrivateKeyStr(String privateKeyStr, byte[] data) {
		return new String(decryptByPrivateKey(toPrivateKey(privateKeyStr), data), StandardCharsets.UTF_8);
	}
	public static String decryptToStrByPrivateKeyStr(String privateKeyStr, String data) {
		return new String(decryptByPrivateKey(toPrivateKey(privateKeyStr), Base64.getDecoder().decode(data)), StandardCharsets.UTF_8);
	}

	// 以下是私钥签名相关方法 ---------- ---------- ---------- ---------- ----------

	@Getter
	@AllArgsConstructor
	public enum RsaSignatureAlgorithmEnum {
		SHA256_WITH_RSA("SHA256withRSA", null),
		SHA512_WITH_RSA("SHA512withRSA", null),
		// 下面两个需要引入 BouncyCastle
		SHA256_WITH_RSA_PSS("SHA512withRSAandMGF1", new PSSParameterSpec("SHA-256", "MGF1", MGF1ParameterSpec.SHA256, 32, 1)),
		SHA512_WITH_RSA_PSS("SHA256withRSAandMGF1", new PSSParameterSpec("SHA-512", "MGF1", MGF1ParameterSpec.SHA512, 64, 1));
		private final String algorithm;
		private final PSSParameterSpec pssParameterSpec;
		public boolean isPss() {
			return pssParameterSpec != null;
		}
	}

	/**
	 * 私钥签名
	 */
	public static byte[] sign(PrivateKey privateKey, RsaSignatureAlgorithmEnum rsaSignatureAlgorithm, byte[] data) {
		try {
			Signature signature = Signature.getInstance(rsaSignatureAlgorithm.getAlgorithm());
			if (rsaSignatureAlgorithm.isPss()) {
				signature.setParameter(rsaSignatureAlgorithm.getPssParameterSpec());
			}
			signature.initSign(privateKey);
			signature.update(data);
			return signature.sign();
		} catch (Throwable cause) {
			throw new RuntimeException(cause);
		}
	}
	public static String signToStr(PrivateKey privateKey, byte[] data) {
		return Base64.getEncoder().encodeToString(sign(privateKey, RsaSignatureAlgorithmEnum.SHA256_WITH_RSA, data));
	}
	public static String signToStr(PrivateKey privateKey, String data) {
		return signToStr(privateKey, data.getBytes(StandardCharsets.UTF_8));
	}
	public static String signToStr(String privateKeyStr, byte[] data) {
		return signToStr(toPrivateKey(privateKeyStr), data);
	}
	public static String signToStr(String privateKeyStr, String data) {
		return signToStr(toPrivateKey(privateKeyStr), data.getBytes(StandardCharsets.UTF_8));
	}

	// 以下是公钥验签相关方法 ---------- ---------- ---------- ---------- ----------

	public static boolean verify(PublicKey publicKey, RsaSignatureAlgorithmEnum rsaSignatureAlgorithm, byte[] data, String sign) {
		try {
			Signature signature = Signature.getInstance(rsaSignatureAlgorithm.getAlgorithm());
			signature.initVerify(publicKey);
			signature.update(data);
			return signature.verify(Base64.getDecoder().decode(sign));
		} catch (Throwable cause) {
			throw new RuntimeException(cause);
		}
	}

	public static boolean verify(PublicKey publicKey, byte[] data, String sign) {
		return verify(publicKey, RsaSignatureAlgorithmEnum.SHA256_WITH_RSA, data, sign);
	}

	public static boolean verify(PublicKey publicKey, String data, String sign) {
		return verify(publicKey, data.getBytes(StandardCharsets.UTF_8), sign);
	}

	public static boolean verify(String publicKeyStr, byte[] data, String sign) {
		return verify(toPublicKey(publicKeyStr), data, sign);
	}

	public static boolean verify(String publicKeyStr, String data, String sign) {
		return verify(toPublicKey(publicKeyStr), data.getBytes(StandardCharsets.UTF_8), sign);
	}

	public static void main(String[] args) {
		KeyPair keyPair = generateKeyPair();
		PrivateKey privateKey = keyPair.getPrivate();
		PublicKey publicKey = keyPair.getPublic();
		String publicKeyStr = toKeyStr(publicKey);
		String privateKeyStr = toKeyStr(privateKey);
		log.info("PublicKey: {}", publicKeyStr);
		log.info("PrivateKey: {}", privateKeyStr);
		String string = "你好";

		String encryptedStr = encryptToStrByPublicKey(publicKey, string);
		log.info("Encrypted: {}", encryptedStr);
		String decryptedStr = decryptToStrByPrivateKey(privateKey, encryptedStr);
		log.info("Decrypted: {}", decryptedStr);

		encryptedStr = encryptToStrByPublicKey(publicKey, string.getBytes(StandardCharsets.UTF_8));
		log.info("Encrypted: {}", encryptedStr);
		decryptedStr = decryptToStrByPrivateKey(privateKey, encryptedStr);
		log.info("Decrypted: {}", decryptedStr);

		encryptedStr = encryptToStrByPublicKeyStr(publicKeyStr, string);
		log.info("Encrypted: {}", encryptedStr);
		decryptedStr = decryptToStrByPrivateKey(privateKey, encryptedStr);
		log.info("Decrypted: {}", decryptedStr);

		encryptedStr = encryptToStrByPublicKeyStr(publicKeyStr, string.getBytes(StandardCharsets.UTF_8));
		log.info("Encrypted: {}", encryptedStr);
		decryptedStr = decryptToStrByPrivateKey(privateKey, encryptedStr);
		log.info("Decrypted: {}", decryptedStr);

		String sign = signToStr(privateKey, string);
		log.info("Sign: {}", sign);
		log.info("Verify: {}", verify(publicKey, string, sign));
		log.info("Verify: {}", verify(publicKey, string.getBytes(StandardCharsets.UTF_8), sign));
		log.info("Verify: {}", verify(publicKeyStr, string, sign));
		log.info("Verify: {}", verify(publicKeyStr, string.getBytes(StandardCharsets.UTF_8), sign));
	}

}
